/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package com.aliasi.crf;

import com.aliasi.chunk.Chunker;
import com.aliasi.chunk.Chunking;
import com.aliasi.chunk.TagChunkCodec;
import com.aliasi.corpus.Corpus;
import com.aliasi.corpus.ObjectHandler;
import com.aliasi.crf.ChainCrfChunker.PreTagging;
import com.aliasi.io.Reporter;
import com.aliasi.io.Reporters;
import com.aliasi.stats.AnnealingSchedule;
import com.aliasi.stats.RegressionPrior;
import com.aliasi.tag.StringTagging;
import com.aliasi.tag.Tagging;
import com.aliasi.tokenizer.Tokenizer;
import com.aliasi.tokenizer.TokenizerFactory;
import com.aliasi.util.AbstractExternalizable;
import com.aliasi.util.Strings;
import entityopinions.TokenWithLEX;
import java.io.IOException;
import java.io.ObjectInput;
import java.io.ObjectOutput;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 *
 * @author Diogo
 */
public class ChainCrfChunkerLEX implements Chunker, Serializable {

    private final ChainCrf mCrf;
    private final TagChunkCodec mCodec;
    private final TokenizerFactory mTokenizerFactory;

    public ChainCrfChunkerLEX(ChainCrf crf,
            TokenizerFactory tokenizerFactory,
            TagChunkCodec codec) {
        mTokenizerFactory = tokenizerFactory;
        mCrf = crf;
        mCodec = codec;
    }

    public ChainCrf ChainCrf() {
        return mCrf;
    }

    public Chunking chunk(CharSequence cSeq) {
        if (cSeq instanceof TokenWithLEX) {
            char[] cs = Strings.toCharArray(cSeq);
            ChainCrfChunker.PreTagging preTagging = preTag(cs, 0, cs.length);
            List<String> tokens_str = preTagging.mTokens;

            List<Integer> lexResults = ((TokenWithLEX) cSeq).getLEXResults();
            List<TokenWithLEX> tokens = new ArrayList<TokenWithLEX>();

            for (int i = 0; i < tokens_str.size(); i++) {
                String str = tokens_str.get(i);
                TokenWithLEX elem = new TokenWithLEX(str, lexResults.get(i));
                tokens.add(elem);
            }
            Tagging tagging = mCrf.tag(tokens);
            return ChainCrfChunkerLEX.toChunking(tagging, preTagging, cs, 0, cs.length, mCodec);
        } else {
            char[] cs = Strings.toCharArray(cSeq);
            return chunk(cs, 0, cs.length);
        }
    }

    static Chunking toChunking(Tagging<String> tagging, PreTagging preTagging,
            char[] cs, int start, int end, TagChunkCodec codec) {

        String s = new String(cs, start, end - start);
        List<String> tokens = preTagging.mTokens;
        int[] tokenStarts = preTagging.mTokenStarts;
        int[] tokenEnds = preTagging.mTokenEnds;
        List<String> tags = tagging.tags();

        StringTagging stringTagging = new StringTagging(tokens, tags, s, tokenStarts, tokenEnds);
        return codec.toChunking(stringTagging);
    }

    public Chunking chunk(char[] cs, int start, int end) {
        ChainCrfChunker.PreTagging preTagging = preTag(cs, start, end);
        List<String> tokens = preTagging.mTokens;
        Tagging<String> tagging = mCrf.tag(tokens);
        return ChainCrfChunker.toChunking(tagging, preTagging, cs, start, end, mCodec);
    }

    ChainCrfChunker.PreTagging preTag(char[] cs, int start, int end) {
        List<Integer> tokenStarts = new ArrayList<Integer>();
        List<Integer> tokenEnds = new ArrayList<Integer>();
        List<String> tokens = new ArrayList<String>();
        Tokenizer tokenizer = mTokenizerFactory.tokenizer(cs, start, end - start);
        String token;
        while ((token = tokenizer.nextToken()) != null) {
            tokens.add(token);
            tokenStarts.add(tokenizer.lastTokenStartPosition());
            tokenEnds.add(tokenizer.lastTokenEndPosition());
        }
        return new ChainCrfChunker.PreTagging(tokens,
                ChainCrfChunker.toArray(tokenStarts),
                ChainCrfChunker.toArray(tokenEnds));
    }

    public static ChainCrfChunkerLEX estimate(Corpus<ObjectHandler<Chunking>> chunkingCorpus,
            TagChunkCodec codec,
            TokenizerFactory tokenizerFactory,
            ChainCrfFeatureExtractor<String> featureExtractor,
            boolean addInterceptFeature,
            int minFeatureCount,
            boolean cacheFeatureVectors,
            RegressionPrior prior,
            int priorBlockSize,
            AnnealingSchedule annealingSchedule,
            double minImprovement,
            int minEpochs,
            int maxEpochs,
            Reporter reporter) throws IOException {
        if (reporter == null) {
            reporter = Reporters.silent();
        }
        reporter.info("Training chain CRF chunker");
        reporter.info("Converting chunk corpus to tag corpus using codec.");
        Corpus<ObjectHandler<Tagging<String>>> taggingCorpus = new ChainCrfChunker.TagCorpus(chunkingCorpus, codec);
        ChainCrf crf = ChainCrf.estimate(taggingCorpus,
                featureExtractor,
                addInterceptFeature,
                minFeatureCount,
                cacheFeatureVectors,
                ChainCrfChunker.ALLOW_UNSEEN_TAG_TRANSITIONS,
                prior,
                priorBlockSize,
                annealingSchedule,
                minImprovement,
                minEpochs,
                maxEpochs,
                reporter);
        return new ChainCrfChunkerLEX(crf, tokenizerFactory, codec);
    }

    static class Serializer extends AbstractExternalizable {

        static final long serialVersionUID = 2460314741682974199L;
        private final ChainCrfChunkerLEX mChunker;

        public Serializer() {
            this(null);
        }

        public Serializer(ChainCrfChunkerLEX chunker) {
            mChunker = chunker;
        }

        public Object read(ObjectInput in) throws IOException, ClassNotFoundException {
            @SuppressWarnings("unchecked")
            ChainCrf<String> crf = (ChainCrf<String>) in.readObject();
            @SuppressWarnings("unchecked")
            TokenizerFactory factory = (TokenizerFactory) in.readObject();
            @SuppressWarnings("unchecked")
            TagChunkCodec codec = (TagChunkCodec) in.readObject();
            return new ChainCrfChunkerLEX(crf, factory, codec);

        }

        public void writeExternal(ObjectOutput out) throws IOException {
            out.writeObject(mChunker.mCrf);
            out.writeObject(mChunker.mTokenizerFactory);
            out.writeObject(mChunker.mCodec);
        }
    }
}
